package algorithms;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;

import framework.State;

public class QLearning {
	private static double gamma = 0.9, errorP = 0.1, epsilon = 0.9;
	private static Map<String, String> policy = new HashMap<String, String>();
	private static Map<String, Double> qvalue = new HashMap<String, Double>();
	private static Map<String, Integer> visits = new HashMap<String, Integer>();
	private static int numberOfRepetitions = 1000;
	private static double[] goalState = new double[] { 2.0, 2.0, 1.0, 2.0, 0.0 };
	private static double learningRate = 1;
	private static double alpha = 1.;
	private static Random rnd = new Random();

	/**
	 * 
	 */
	public static void init() {
		// Set 0 Q-value for every state
		for (String s : State.allStates.keySet()) {
			qvalue.put(s, 0.);
			visits.put(s, 0);
		}

		// Set a qvalue of 100 for the goal state
		State.allStates.get(State.toString(goalState))[4] = 100;
	}

	/**
	 * 
	 */
	public static void iterate() {
		int visitCount = 0;
		for (int i = 0; i < numberOfRepetitions; i++) {
			// Start with a random state
			Object[] values = State.allStates.values().toArray();
			double[] state = (double[]) values[rnd.nextInt(values.length)];
			visitCount = visits.get(State.toString(state)) + 1;
			visits.put(State.toString(state), visitCount);
			while (!State.toString(state).equals(State.toString(goalState))) {
				// Select an action determined by the selection strategy.
				double[] nextState = select(state);
				visitCount = visits.get(State.toString(nextState)) + 1;
				visits.put(State.toString(nextState), visitCount);

				// Using this possible action, consider to go to the next state
				double q = qvalue.get(State.toString(state));
				double maxQ = maxQ(nextState);
				double r = State.getValue(nextState);

				// Set learningRate
				//learningRate = 1 / visits.get(State.toString(nextState));
				Integer visitss = visits.get(State.toString(nextState));
				learningRate = Math.pow((double)visitss, alpha * -1);
				
				double value = q + learningRate * (r + gamma * maxQ - q);
				qvalue.put(State.toString(state), value);
				state[4] = value;

				// Set the next state as the current state
				state = nextState;
			}
		}
		// Update policy before printing it.
		policy.clear();
		for (String s : State.allStates.keySet()) {
			policy.put(s, State.toString(maxQState(State.allStates.get(s))));
		}
	}

	/**
	 * 
	 * @param curState
	 * @return
	 */
	public static double[] select(double[] state) {
		// Epsilon-greedy with possible errors
		double[] nextState = null;

		if (rnd.nextDouble() < epsilon) {
			// Random exploration
			List<String> states = State.getPossibleTransitions(state);
			if (states.size() == 0) {
				return null;
			}
			String s2 = states.get(rnd.nextInt(states.size()));
			nextState = State.allStates.get(s2);
		} else {
			// Exploitation
			nextState = maxQState(state);
		}

		// Error
		if (rnd.nextDouble() < errorP) {
			// Find out which disk we are moving
			int move;
			if (nextState[0] != state[0]) {
				move = 0; // A Moved
			} else {
				move = 1; // B Moved
			}

			// Get the resulting state from the error
			for (String s : State.getPossibleTransitions(nextState)) {
				double[] errorState = State.allStates.get(s);
				if (errorState[move] != state[move] && errorState[move] != nextState[move]) {
					nextState = errorState;
				}
			}
		}

		return nextState;
	}

	/**
	 * 
	 * @param s
	 * @return
	 */
	private static double maxQ(double[] state) {
		List<String> states = State.getPossibleTransitions(state);
		if (states.size() == 0) {
			return -1;
		}
		double maxValue = Double.MIN_VALUE;
		for (String nextState : states) {
			double value = qvalue.get(nextState);
			if (value > maxValue)
				maxValue = value;
		}
		return maxValue;
	}

	/**
	 * 
	 * @param state
	 * @return
	 */
	private static double[] maxQState(double[] state) {
		List<String> states = State.getPossibleTransitions(state);
		if (states.size() == 0) {
			return state;
		}
		double maxValue = Double.MIN_VALUE;
		String maxState = null;
		for (String nextState : states) {
			double value = qvalue.get(nextState);
			if (value > maxValue)
				maxValue = value;
			maxState = nextState;
		}
		return State.allStates.get(maxState);
	}

	/**
	 * 
	 * @return
	 */
	public static String printPolicy() {
		StringBuilder builder = new StringBuilder();

		builder.append("----Policy-----\n");
		for (double[] s : State.allStates.values()) {
			// builder.append(State.toString(s));
			builder.append(State.toVisualString(s, true));
			builder.append("-->\n");
			// builder.append(policy.get(State.toString(s)));
			builder.append(State.toVisualString(State.allStates.get(policy.get(State.toString(s))),
					true));
			builder.append("########\n");
		}
		return builder.toString();
	}
}
